use std::fs;

use crate::Puzzle;

type Bits = Vec<bool>;

pub struct BinaryDiagnostic {
    bits_list: Vec<Bits>,
    bits_size: usize,
}

impl BinaryDiagnostic {
    #[allow(dead_code)]
    fn new_with_list(list: Vec<Bits>) -> Self {
        Self {
            bits_size: list.first().unwrap_or(&vec![]).len(),
            bits_list: list,
        }
    }

    fn count_one_and_zero(&self) -> (Vec<u32>, Vec<u32>) {
        let length = self.bits_size;

        let mut ones: Vec<u32> = vec![0; length];
        let mut zeros: Vec<u32> = vec![0; length];

        for bits in &self.bits_list {
            for (index, bit) in bits.iter().enumerate() {
                if *bit {
                    ones[index] += 1;
                } else {
                    zeros[index] += 1;
                }
            }
        }

        (ones, zeros)
    }

    fn count_one_and_zero_at_index(&self, lists: &Vec<Bits>, index: usize) -> (u32, u32) {
        let mut ones_count: u32 = 0;
        let mut zeros_count: u32 = 0;

        for bits in lists {
            if let Some(bit) = bits.get(index) {
                if *bit {
                    ones_count += 1;
                } else {
                    zeros_count += 1;
                }
            }
        }

        (ones_count, zeros_count)
    }

    fn bits_to_decimal(bits: &Bits) -> u32 {
        bits.iter()
            .rev()
            .enumerate()
            .fold(0, |result, (index, value)| {
                if *value {
                    result + 2u32.pow(index as u32)
                } else {
                    result
                }
            })
    }
}

impl Puzzle for BinaryDiagnostic {
    type Answer = u32;

    fn new() -> Self {
        let input = fs::read_to_string("input/day03.md").unwrap();

        let list = input
            .lines()
            .map(|line| line.trim().chars().map(|c| c == '1').collect::<Bits>())
            .filter(|line| !line.is_empty())
            .collect::<Vec<Bits>>();

        Self {
            bits_size: list.first().unwrap_or(&vec![]).len(),
            bits_list: list,
        }
    }

    fn solve_part1(&self) -> Self::Answer {
        let (ones, zeros) = self.count_one_and_zero();

        let value1 = ones
            .iter()
            .zip(zeros.iter())
            .map(|(c1, c0)| c1 > c0)
            .collect::<Vec<bool>>();
        let value2 = value1.iter().map(|v| !v).collect::<Vec<bool>>();

        let value1 = Self::bits_to_decimal(&value1);
        let value2 = Self::bits_to_decimal(&value2);

        value1 * value2
    }

    fn solve_part2(&self) -> Self::Answer {
        let mut list1 = self.bits_list.to_vec();
        for index in 0..self.bits_size {
            let (ones_count, zeros_count) = self.count_one_and_zero_at_index(&list1, index);
            let target_bit = ones_count >= zeros_count;

            list1 = list1
                .into_iter()
                .filter(|bits| bits[index] == target_bit)
                .collect::<Vec<Bits>>();

            if list1.len() == 1 {
                break;
            }
        }

        let mut list2 = self.bits_list.to_vec();
        for index in 0..self.bits_size {
            let (ones_count, zeros_count) = self.count_one_and_zero_at_index(&list2, index);
            let target_bit = ones_count < zeros_count;

            list2 = list2
                .into_iter()
                .filter(|bits| bits[index] == target_bit)
                .collect::<Vec<Bits>>();

            if list2.len() == 1 {
                break;
            }
        }

        let Some(bits1) = list1.first() else { return 0 };
        let Some(bits2) = list2.first() else { return 0 };
        let value1 = Self::bits_to_decimal(bits1);
        let value2 = Self::bits_to_decimal(bits2);
        value1 * value2
    }
}

#[cfg(test)]
mod test {
    use super::*;

    fn create_list() -> Vec<Vec<bool>> {
        let raw = "
            00100
            11110
            10110
            10111
            10101
            01111
            00111
            11100
            10000
            11001
            00010
            01010
        ";

        raw.lines()
            .map(|line| line.trim().chars().map(|c| c == '1').collect::<Vec<bool>>())
            .filter(|line| !line.is_empty())
            .collect::<Vec<Vec<bool>>>()
    }

    #[test]
    fn part1() {
        let list = create_list();
        assert_eq!(BinaryDiagnostic::new_with_list(list).solve_part1(), 198);
    }

    #[test]
    fn part2() {
        let list = create_list();
        assert_eq!(BinaryDiagnostic::new_with_list(list).solve_part2(), 230);
    }
}
